from .image_convert_base import ConvertBase
import cv2 as cv
import numpy as np

# 随机贴图覆盖
class LocalOcclusionConvert(ConvertBase):
    
    def __init__(self, min_size, max_size, use_rate = 0.5, fill_color = [0, 0, 0]):
        super().__init__(use_rate)
        
        self.min_size = min_size
        self.max_size = max_size
        self.fill_color = fill_color
    
    def __get_occlusion(self, img):
        
        img_h, img_w, _ = img.shape
        
        w = np.random.randint(self.min_size, self.max_size)
        h = np.random.randint(self.min_size, self.max_size)
        
        start_x = np.random.randint(0, img_w - w)
        start_y = np.random.randint(0, img_h - h)
        
        return [start_x, start_x + w, start_y, start_y + h]
    
    
    
    def convert(self, img, boxes, points):
        
        s_x, w, s_y, h = self.__get_occlusion(img)
        
        img[h:s_y, w:s_x] = self.fill_color
        
        return img, boxes, points
        